# This kit is for testing the generated model, check the accuracy of the labeled objects vs the
# given labeled image which used for training the model.
import io
import os
import sys
import shutil

import PIL.Image as Image
import numpy as np
import tensorflow as tf

sys.path.append("..")
from object_detection.utils import label_map_util
from tensorflow.python.lib.io.file_io import FileIO

import tfrecord_decoder
from tools.utils import predict_result_utils as utils

# Change the path to the utils package under models/research/object_detection.
sys.path.append('../../models/research/object_detection/utils')
import visualization_utils


flags = tf.app.flags

flags.DEFINE_string("model_path", 'G:/to_be_test/frozen_inference_graph_leo_with_or_without_helmet.pb',
                    "The model used for training.")
flags.DEFINE_string('path_to_labels',
                    'G:/to_be_test/label.pbtxt',
                    "The path for label mapping .pbtxt file")
flags.DEFINE_string('tf_record_path', 'G:/to_be_test/test.record',
                    "The path to the config file.")
flags.DEFINE_integer('epoch', 10,
                     'The number of samples to test in a single run.')
flags.DEFINE_integer('test_buffer_size', 20,
                     'The buffer size for testing within the TF record.')
FLAGS = flags.FLAGS

PATH_TO_CKPT = FLAGS.model_path
TFRECORD_FILE = FLAGS.tf_record_path
TESTING_EPOCH = FLAGS.epoch
BUFFER_SIZE = FLAGS.test_buffer_size

NUM_CLASSES = 90


# The accuracy is calculated with the cross area of the generated label with the given area.
def calculate_accuracy(groundtruth_boxes, normalized_boxes, groundtruth_text, texts, iouThreshold=0.8):
    image_accuracy = []
    numOffound = 0
    for i in range(groundtruth_boxes.__len__()):
        max_accuracy = -1
        label_box = groundtruth_boxes[i]
        text = groundtruth_text[i].decode('utf-8')
        ymin, xmin, ymax, xmax = label_box[0]
        xmin = float(xmin)
        ymin = float(ymin)
        xmax = float(xmax)
        ymax = float(ymax)
        groundtruth_area = (xmax - xmin) * (ymax - ymin)
        for j in range(normalized_boxes.__len__()):
            test_box = normalized_boxes[j]
            test_text = str(texts[j][0])
            test_text = test_text[:test_text.index(':')]
            t_ymin, t_xmin, t_ymax, t_xmax = test_box[0]

            tarea = (t_xmax - t_xmin) * (t_ymax - t_ymin)

            x1 = max(xmin, t_xmin)
            x2 = min(xmax, t_xmax)
            y1 = max(ymin, t_ymin)
            y2 = min(ymax, t_ymax)

            w = max(0, x2 - x1)
            h = max(0, y2 - y1)

            area = w * h
            iou = area / (groundtruth_area + tarea - area)

            if iou > max_accuracy:
                max_accuracy = iou

        if max_accuracy > iouThreshold:
            numOffound = numOffound + 1

        if max_accuracy > 0:
            image_accuracy.append(max_accuracy)
        else:
            image_accuracy.append(0)

        numOfGroundTruth = groundtruth_boxes.__len__()
        numOfPredicted = normalized_boxes.__len__()

    return (numOfGroundTruth, numOfPredicted, numOffound, 0) if image_accuracy.__len__() == 0 else (numOfGroundTruth, numOfPredicted, numOffound, sum(image_accuracy) / image_accuracy.__len__())



def calculate_average_accuracy(total_accuracy_indices):
    return sum(total_accuracy_indices) / total_accuracy_indices.__len__()


def process_accuracy(filename, height, width, groundtruth_boxes, groundtruth_text, box_to_display_map):
    normalized_boxes, texts = convert_to_normalized_boxes(box_to_display_map, float(height), float(width))
    normalized_groundtruth_boxes = convert_box_values(groundtruth_boxes, float(height), float(width))

    numOfGroundTruth, numOfPredicted, numOffound, single_img_accuracy = calculate_accuracy(normalized_groundtruth_boxes, normalized_boxes, groundtruth_text, texts)
    print("File Name:" + filename.decode('utf-8')[:-4] 
        + "\tnumOfGroundTruth:" + str(numOfGroundTruth) 
        + "\tnumOfPredicted:" + str(numOfPredicted)
        + "\tnumOffound:" + str(numOffound)
        + "\tAccuracy:" + str(single_img_accuracy))
    return numOfGroundTruth, numOfPredicted, numOffound, single_img_accuracy


def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape(
        (im_height, im_width, 3)).astype(np.uint8)


def draw_boxes_on_image_array(image, groundtruth_boxes, groundtruth_text):
    texts = []
    for text in groundtruth_text:
        text = text.decode('utf-8')
        texts.append(text)
    for i in range(groundtruth_boxes.__len__()):
        ymin, xmin, ymax, xmax = groundtruth_boxes[i]
        image_pil = Image.fromarray(image).convert('RGB')
        text_arr = []
        text_arr.append(texts[i])
        visualization_utils.draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax,
                                                       color='red', thickness=4,
                                                       display_str_list=text_arr, use_normalized_coordinates=True)
        np.copyto(image, np.array(image_pil))


def process_testing(tfrecord_file, buffer_size, epoch, path_to_ckpt, path_to_labels):
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with FileIO(path_to_ckpt, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
            label_map = label_map_util.load_labelmap(path_to_labels)
            categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
                                                                        use_display_name=True)
            category_index = label_map_util.create_category_index(categories)

    with detection_graph.as_default():
        dataset = tf.data.TFRecordDataset(tfrecord_file)
        dataset = dataset.shuffle(buffer_size=buffer_size)
        iterator = dataset.make_one_shot_iterator()
        element = iterator.get_next()
        with tf.Session(graph=detection_graph) as sess:
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            total_accuracy = []
            total_numOfGroundTruth = []
            total_numOfPredicted = []
            total_numOffound = []

            for i in range(epoch):
                example = sess.run(element)
                image, height, width, filename, xmin, xmax, ymin, ymax, text \
                    = tfrecord_decoder.decode_example(example)
                e_image, e_height, e_width, e_filename, e_xmin, e_xmax, e_ymin, e_ymax, e_text = \
                    sess.run([image, height, width, filename, xmin, xmax, ymin, ymax, text])

                groundtruth_boxes, groundtruth_text = \
                    convert_tensor_dataset_to_box_and_text(e_ymin, e_xmin, e_ymax, e_xmax, e_text)

                image_io = io.BytesIO(e_image)
                image_data = Image.open(image_io)

                # image_byte = load_image_into_numpy_array(image_data)
                # image_np_expanded = np.expand_dims(image_byte, axis=0)

                # placeholder for single channel images
                image_np = load_image_into_numpy_array(image_data)
                image_np_original = load_image_into_numpy_array(image_data)
                image_np1 = np.concatenate((image_np, image_np), axis=2)
                image_np = np.concatenate((image_np1, image_np), axis=2)

                # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
                image_np_expanded = np.expand_dims(image_np_original, axis=0)
                image_np_expanded1 = np.concatenate((image_np_expanded, image_np_expanded), axis=3)
                image_np_expanded = np.concatenate((image_np_expanded1, image_np_expanded), axis=3)

                # Do the prediction and returns the label boxes.
                (boxes, scores, classes, num) = sess.run(
                    [detection_boxes, detection_scores, detection_classes, num_detections],
                    feed_dict={image_tensor: image_np_expanded})

                box_to_display_str_map = visualization_utils.visualize_boxes_and_labels_on_image_array(
                    image_np,
                    np.squeeze(boxes),
                    np.squeeze(classes).astype(np.int32),
                    np.squeeze(scores),
                    category_index,
                    use_normalized_coordinates=True,
                    line_thickness=8)

                draw_boxes_on_image_array(image_np, groundtruth_boxes, groundtruth_text)

                # Calculate the accuracy.
                numOfGroundTruth, numOfPredicted, numOffound, accuracy = process_accuracy(e_filename, e_height, e_width, groundtruth_boxes, groundtruth_text,
                                            box_to_display_str_map)
                total_accuracy.append(accuracy)
                total_numOfGroundTruth.append(numOfGroundTruth)
                total_numOfPredicted.append(numOfPredicted)
                total_numOffound.append(numOffound)

                result = Image.fromarray(image_np)
                result_file_name = e_filename.decode('utf-8')[:-4] + '_acc_' + str(accuracy) + '.jpeg'
                io.FileIO(result_file_name, 'w')
                result.save(result_file_name, 'jpeg')
                # os.rename(result_file_name, '/object_tutorial/result/' + result_file_name)
                shutil.move(result_file_name, './object_tutorial/result/' + result_file_name)

    accuracy_sum = sum(total_accuracy)
    numOfGroundTruth_sum = sum(total_numOfGroundTruth)
    numOfPredicted_sum = sum(total_numOfPredicted)
    numOffound_sum = sum(total_numOffound)

    print("=========================================================")
    print("Total Average Accuracy:" + str(accuracy_sum / total_accuracy.__len__() if total_accuracy.__len__() > 0 else 0))
    print("Total Number of GroundTruth:" + str(numOfGroundTruth_sum))
    print("Total Number of Predicted:" + str(numOfPredicted_sum))
    print("Total Number of found:" + str(numOffound_sum))


def main():
    process_testing(TFRECORD_FILE,
                    BUFFER_SIZE,
                    TESTING_EPOCH,
                    PATH_TO_CKPT,
                    FLAGS.path_to_labels)


if __name__ == '__main__':
    main()
